import os
import time
import cv2
import argparse
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from PIL import Image, ImageTk
from ultralytics import YOLO
from collections import Counter

# 定义推理参数
class Config:
    # 模型相关参数
    weights = r'C:\Users\86187\Desktop\optimal_knn_mnist_question_2-main\口罩佩戴识别\runs\train\weights\best.pt'  # 训练好的模型权重
    
    # 推理相关参数
    source = 'data/prepared/test/images'    # 推理图像或视频源
    img_size = 640                          # 图像大小
    conf_thres = 0.25                       # 置信度阈值
    iou_thres = 0.45                        # IoU阈值
    
    # 输出相关参数
    save = True                             # 是否保存结果 
    save_txt = False                        # 是否保存文本结果
    save_conf = False                       # 是否保存置信度到文本文件
    show = True                             # 是否显示结果
    
    # 其他配置
    device = ''                             # 设备选择 ('cpu' or '0,1,2,3')
    classes = None                          # 过滤特定类别，例如 [0, 1] 或 None
    agnostic_nms = False                    # 类别无关的NMS
    augment = False                         # 增强推理
    project = 'runs'                        # 保存项目路径
    name = 'detect'                         # 检测名称
    exist_ok = True                         # 是否允许覆盖现有实验目录
    line_thickness = 3                      # 绘制边界框的线宽
    hide_labels = False                     # 是否隐藏标签
    hide_conf = False                       # 是否隐藏置信度

def detect_image(weights, source, conf_thres=0.25, iou_thres=0.45, img_size=640, 
                 save=True, show=True, device='', project='runs', name='detect'):
    """
    使用训练好的模型进行图像目标检测
    """
    print(f"Loading model from {weights}...")
    model = YOLO(weights)
    
    print(f"Running inference on {source}...")
    results = model.predict(
        source=source,
        conf=conf_thres,
        iou=conf_thres,
        imgsz=img_size,
        save=save,
        show=show,
        device=device,
        project=project,
        name=name,
        exist_ok=True,
        line_thickness=3
    )
    
    # 计算检测到的目标数量
    class_counts = Counter()
    total_objects = 0
    
    for result in results:
        if len(result.boxes) > 0:
            # 获取类别ID和类别名称
            for box in result.boxes:
                cls_id = int(box.cls.item())
                cls_name = result.names[cls_id]
                class_counts[cls_name] += 1
                total_objects += 1
    
    # 输出检测结果摘要
    print("\n--- 检测结果摘要 ---")
    print(f"总共检测到 {total_objects} 个目标")
    for cls_name, count in class_counts.items():
        print(f"{cls_name}: {count} 个")
    
    print(f"\n检测完成。结果保存在 {project}/{name}")
    return results

def detect_webcam(weights, conf_thres=0.25, iou_thres=0.45, img_size=640, 
                 device='', cam_id=0):
    """
    使用训练好的模型进行实时摄像头目标检测
    """
    print(f"Loading model from {weights}...")
    model = YOLO(weights)
    
    # 打开摄像头
    print(f"Opening webcam {cam_id}...")
    cap = cv2.VideoCapture(cam_id)
    
    # 检查摄像头是否成功打开
    if not cap.isOpened():
        print("错误：无法打开摄像头")
        return
    
    print("开始实时检测. 按 'q' 退出...")
    
    while True:
        # 读取一帧
        ret, frame = cap.read()
        if not ret:
            print("无法获取摄像头图像，退出...")
            break
        
        # 进行目标检测
        results = model.predict(
            source=frame,
            conf=conf_thres,
            iou=iou_thres,
            imgsz=img_size,
            device=device,
            stream=True
        )
        
        # 处理结果
        for result in results:
            # 获取带注释的帧
            annotated_frame = result.plot()
            
            # 显示检测结果
            cv2.imshow("YOLOv8 口罩检测", annotated_frame)
        
        # 按q退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()
    print("检测结束")

def create_gui():
    """
    创建GUI界面，支持上传文件和摄像头选择
    """
    # 创建主窗口
    root = tk.Tk()
    root.title("口罩佩戴检测系统")
    root.geometry("800x600")
    
    # 设置样式
    style = ttk.Style()
    style.configure("TButton", font=("微软雅黑", 12))
    style.configure("TLabel", font=("微软雅黑", 12))
    
    # 创建标题标签
    title_label = ttk.Label(root, text="基于YOLOv8的口罩佩戴检测系统", font=("微软雅黑", 18, "bold"))
    title_label.pack(pady=20)
    
    # 创建模型路径输入
    model_frame = ttk.Frame(root)
    model_frame.pack(pady=10, fill="x", padx=50)
    
    model_label = ttk.Label(model_frame, text="模型路径:")
    model_label.pack(side="left", padx=5)
    
    model_path = tk.StringVar(value=Config.weights)
    model_entry = ttk.Entry(model_frame, textvariable=model_path, width=50)
    model_entry.pack(side="left", padx=5, fill="x", expand=True)
    
    def browse_model():
        filename = filedialog.askopenfilename(
            title="选择模型文件",
            filetypes=(("PyTorch模型", "*.pt"), ("所有文件", "*.*"))
        )
        if filename:
            model_path.set(filename)
    
    browse_button = ttk.Button(model_frame, text="浏览...", command=browse_model)
    browse_button.pack(side="left", padx=5)
    
    # 创建阈值设置
    threshold_frame = ttk.Frame(root)
    threshold_frame.pack(pady=10, fill="x", padx=50)
    
    conf_label = ttk.Label(threshold_frame, text="置信度阈值:")
    conf_label.pack(side="left", padx=5)
    
    conf_value = tk.DoubleVar(value=Config.conf_thres)
    conf_scale = ttk.Scale(threshold_frame, from_=0.1, to=1.0, orient="horizontal", 
                           variable=conf_value, length=200)
    conf_scale.pack(side="left", padx=5)
    
    conf_display = ttk.Label(threshold_frame, textvariable=conf_value)
    conf_display.pack(side="left", padx=5)
    
    # 创建检测模式框架
    mode_frame = ttk.LabelFrame(root, text="检测模式", padding=10)
    mode_frame.pack(pady=20, padx=50, fill="x")
    
    # 创建摄像头选择框架
    camera_frame = ttk.Frame(mode_frame)
    camera_frame.pack(pady=10, fill="x")
    
    camera_label = ttk.Label(camera_frame, text="选择摄像头:")
    camera_label.pack(side="left", padx=5)
    
    camera_id = tk.IntVar(value=0)
    camera_combo = ttk.Combobox(camera_frame, textvariable=camera_id, width=5)
    camera_combo['values'] = (0, 1, 2, 3)
    camera_combo.pack(side="left", padx=5)
    
    def start_camera_detection():
        try:
            model_file = model_path.get()
            if not os.path.exists(model_file):
                messagebox.showerror("错误", f"模型文件不存在: {model_file}")
                return
                
            cam_id = camera_id.get()
            conf = conf_value.get()
            
            # 显示加载提示
            status_label.config(text="正在启动摄像头检测...")
            root.update()
            
            # 启动摄像头检测
            detect_webcam(
                weights=model_file,
                conf_thres=conf,
                cam_id=cam_id
            )
            
            status_label.config(text="摄像头检测已完成")
        except Exception as e:
            messagebox.showerror("错误", f"摄像头检测失败: {str(e)}")
            status_label.config(text="检测失败")
    
    camera_button = ttk.Button(camera_frame, text="启动摄像头检测", 
                              command=start_camera_detection)
    camera_button.pack(side="left", padx=20)
    
    # 创建文件上传框架
    file_frame = ttk.Frame(mode_frame)
    file_frame.pack(pady=10, fill="x")
    
    file_label = ttk.Label(file_frame, text="图像文件:")
    file_label.pack(side="left", padx=5)
    
    file_path = tk.StringVar()
    file_entry = ttk.Entry(file_frame, textvariable=file_path, width=50)
    file_entry.pack(side="left", padx=5, fill="x", expand=True)
    
    def browse_file():
        filename = filedialog.askopenfilename(
            title="选择图像文件",
            filetypes=(
                ("图像文件", "*.jpg *.jpeg *.png *.bmp"),
                ("所有文件", "*.*")
            )
        )
        if filename:
            file_path.set(filename)
    
    browse_file_button = ttk.Button(file_frame, text="浏览...", command=browse_file)
    browse_file_button.pack(side="left", padx=5)
    
    def start_file_detection():
        try:
            model_file = model_path.get()
            image_file = file_path.get()
            
            if not os.path.exists(model_file):
                messagebox.showerror("错误", f"模型文件不存在: {model_file}")
                return
                
            if not os.path.exists(image_file):
                messagebox.showerror("错误", f"图像文件不存在: {image_file}")
                return
            
            conf = conf_value.get()
            
            # 显示加载提示
            status_label.config(text="正在进行检测...")
            root.update()
            
            # 启动文件检测
            results = detect_image(
                weights=model_file,
                source=image_file,
                conf_thres=conf,
                show=False
            )
            
            # 获取检测结果并在GUI中显示
            for result in results:
                # 获取带注释的图像
                annotated_img = result.plot()
                
                # 转换为PIL图像来显示
                img_pil = Image.fromarray(cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB))
                
                # 调整大小以适合显示区域
                width, height = img_pil.size
                max_width = 700
                max_height = 500
                
                # 计算调整后的尺寸
                if width > max_width or height > max_height:
                    ratio = min(max_width / width, max_height / height)
                    new_width = int(width * ratio)
                    new_height = int(height * ratio)
                    img_pil = img_pil.resize((new_width, new_height), Image.LANCZOS)
                
                # 显示图像
                img_tk = ImageTk.PhotoImage(img_pil)
                
                # 更新结果标签
                if hasattr(root, 'result_image_label'):
                    root.result_image_label.config(image=img_tk)
                    root.result_image_label.image = img_tk
                else:
                    root.result_image_label = ttk.Label(root, image=img_tk)
                    root.result_image_label.image = img_tk
                    root.result_image_label.pack(pady=10)
                
                # 显示检测的类别统计
                class_counts = Counter()
                total_objects = 0
                
                if len(result.boxes) > 0:
                    for box in result.boxes:
                        cls_id = int(box.cls.item())
                        cls_name = result.names[cls_id]
                        class_counts[cls_name] += 1
                        total_objects += 1
                
                # 更新状态标签
                status_text = f"检测完成，共发现 {total_objects} 个目标 | "
                for cls_name, count in class_counts.items():
                    status_text += f"{cls_name}: {count}个 | "
                
                status_label.config(text=status_text)
            
        except Exception as e:
            messagebox.showerror("错误", f"文件检测失败: {str(e)}")
            status_label.config(text="检测失败")
    
    file_button = ttk.Button(file_frame, text="检测图像", command=start_file_detection)
    file_button.pack(side="left", padx=20)
    
    # 创建状态标签
    status_frame = ttk.Frame(root)
    status_frame.pack(pady=10, side="bottom", fill="x")
    
    status_label = ttk.Label(status_frame, text="就绪", font=("微软雅黑", 10))
    status_label.pack(pady=5)
    
    # 设置退出按钮
    exit_button = ttk.Button(root, text="退出程序", command=root.destroy)
    exit_button.pack(pady=10, side="bottom")
    
    # 启动主循环
    root.mainloop()

def detect():
    """
    使用训练好的模型进行目标检测
    """
    # 获取配置
    cfg = Config()
    
    # 检查模型权重是否存在
    if not os.path.exists(cfg.weights):
        print(f"错误：模型权重文件不存在: {cfg.weights}")
        return
    
    # 启动GUI界面
    create_gui()

if __name__ == "__main__":
    detect() 